{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "cebea951-cce3-47a2-b7a6-152579a8460c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Segmented reduce result: [ 1  5  9 21  0  9]\n"
     ]
    }
   ],
   "source": [
    "# An attempt to recreate studies/cuda-kernels/reducers/awkward_reduce_sum_atomics.py using cccl instead of raw cuda kernels\n",
    "import awkward as ak\n",
    "import cupy as cp\n",
    "import numpy as np\n",
    "\n",
    "from cuda.compute import segmented_reduce\n",
    "\n",
    "def sum_op(a, b):\n",
    "    return a+b\n",
    "\n",
    "awkward_array = ak.Array([[1], [2, 3], [4, 5], [6, 7, 8], [], [9]], backend = 'cuda')\n",
    "input_data = awkward_array.layout.content.data\n",
    "offsets = awkward_array.layout.offsets.data\n",
    "\n",
    "# Prepare the start and end offsets\n",
    "start_o = offsets[:-1]\n",
    "end_o = offsets[1:]\n",
    "\n",
    "# Prepare the output array\n",
    "n_segments = start_o.size\n",
    "output = cp.empty(n_segments, dtype=cp.int32)\n",
    "\n",
    "# Initial value for the reduction\n",
    "h_init = np.array([0], dtype=np.int32)\n",
    "\n",
    "# Perform the segmented reduce\n",
    "segmented_reduce(\n",
    "    input_data, output, start_o, end_o, sum_op, h_init, n_segments\n",
    ")\n",
    "\n",
    "print(f\"Segmented reduce result: {output.get()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6938f06d-0432-461e-bd21-910c885e0628",
   "metadata": {},
   "source": [
    "Success!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c1a84ba1-a411-40af-b5da-d19e3769e9da",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Verify the result.\n",
    "expected_output = cp.asarray([1, 5, 9, 21, 0, 9], dtype=output.dtype)\n",
    "assert (output == expected_output).all()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eef0da43-bd57-40c6-99a5-4f0888e8f2e2",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "7c8bbeb0-3ae0-4efe-a3c7-308712b45569",
   "metadata": {},
   "source": [
    "### Compare time metrics between cuda-kernels and cccl"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44d88ab0-addc-45d5-abf6-c7d7e9ce8952",
   "metadata": {},
   "source": [
    "Using cuda-kernels:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "f8e92fdf-2153-4ada-a36a-a9bbbba6c12b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cupy as cp\n",
    "\n",
    "cuda_kernel = \"\"\"\n",
    "extern \"C\" {\n",
    "    __global__ void awkward_reduce_sum_a(int* toptr, int* fromptr, int* parents, int lenparents, int outlength) {\n",
    "       int thread_id = blockIdx.x * blockDim.x + threadIdx.x;\n",
    "\n",
    "       if (thread_id < outlength) {\n",
    "          toptr[thread_id] = 0;\n",
    "       }\n",
    "    }\n",
    "}\n",
    "extern \"C\" {\n",
    "    __global__ void awkward_reduce_sum_b(int* toptr, int* fromptr, int* parents, int lenparents, int outlength) {\n",
    "       int thread_id = blockIdx.x * blockDim.x + threadIdx.x;\n",
    "       int stride = blockDim.x * gridDim.x;\n",
    "\n",
    "       for (int i = thread_id; i < lenparents; i += stride) {\n",
    "           atomicAdd(&toptr[parents[i]], fromptr[i]);\n",
    "       }\n",
    "    }\n",
    "}\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "074b68de-4392-4311-bc42-ff4dd0631d11",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "331 μs ± 10.8 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "parents = cp.array([0, 1, 1, 2, 2, 3, 3, 3, 5], dtype=cp.int32)\n",
    "fromptr = cp.array([1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=cp.int32)\n",
    "lenparents = len(parents)\n",
    "outlength = int(cp.max(parents)) + 1\n",
    "toptr = cp.zeros(outlength, dtype=cp.int32)\n",
    "\n",
    "block_size = 256\n",
    "grid_size = (lenparents + block_size - 1) // block_size\n",
    "\n",
    "raw_module = cp.RawModule(code=cuda_kernel)\n",
    "\n",
    "awkward_reduce_sum_a = raw_module.get_function('awkward_reduce_sum_a')\n",
    "awkward_reduce_sum_b = raw_module.get_function('awkward_reduce_sum_b')\n",
    "\n",
    "awkward_reduce_sum_a((grid_size,), (block_size,), (toptr, fromptr, parents, lenparents, outlength))\n",
    "awkward_reduce_sum_b((grid_size,), (block_size,), (toptr, fromptr, parents, lenparents, outlength))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "469dc734-0416-4aec-b366-a5dcd3ce8228",
   "metadata": {},
   "source": [
    "Using cccl:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "id": "904b5c00-40b1-4570-87ee-208171baaaf7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The slowest run took 7.03 times longer than the fastest. This could mean that an intermediate result is being cached.\n",
      "138 μs ± 122 μs per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "awkward_array = rand_arr\n",
    "input_data = awkward_array.layout.content.data \n",
    "offsets = awkward_array.layout.offsets.data\n",
    "\n",
    "# Prepare the start and end offsets\n",
    "start_o = offsets[:-1]\n",
    "end_o = offsets[1:]\n",
    "\n",
    "# Prepare the output array\n",
    "n_segments = start_o.size\n",
    "output = cp.empty(n_segments, dtype=np.int32)\n",
    "\n",
    "# Initial value for the reduction\n",
    "h_init = np.array([0], dtype=np.int32)\n",
    "\n",
    "# Perform the segmented reduce\n",
    "segmented_reduce(\n",
    "    input_data, output, start_o, end_o, sum_op, h_init, n_segments\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb68494f-3345-4959-9457-73767bf7cf54",
   "metadata": {},
   "source": [
    "### Generating a big (2GB) array for testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d19a1a18-b984-41a4-9706-3ca41f9c3075",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre>[[0.5446351700822585, 0.31547762591056666, ..., 0.5086729539863157],\n",
       " [0.7875010981288579],\n",
       " [0.8544666454677607, 0.9858165463059176, ..., 0.48755491796298195],\n",
       " [0.2457779153558392, 0.3902446871547895, ..., 0.7423458178078217],\n",
       " [],\n",
       " [0.19331265794192048, 0.14222699239455688],\n",
       " [0.8229611054203183, 0.839945724915246, 0.7088897522743696],\n",
       " [0.6186932546777535, 0.5898080263100203, ..., 0.7492442360444409],\n",
       " [0.18470051000056348, 0.7886784755850496, ..., 0.7143500962029918],\n",
       " [0.27022460572082646, 0.8581594760243878, ..., 0.6341535146163997],\n",
       " ...,\n",
       " [0.9477070629803681, 0.14588253817215951, 0.9086179064549229],\n",
       " [0.5219715451501898, 0.4417330617577138],\n",
       " [0.08706698485919068, 0.2568640154922687],\n",
       " [0.039257266672439506, 0.38848707930766774],\n",
       " [0.5342610161868544, 0.5172952861006859, 0.6712729045351582],\n",
       " [0.5009163220102286, 0.17350017668502699, ..., 0.03325255729962834],\n",
       " [0.9770648067429579, 0.2153223601906204, ..., 0.9051232112468341],\n",
       " [0.9547518307181262],\n",
       " [0.8511949108999146, 0.32843391498125457, ..., 0.12210666613637128]]\n",
       "---------------------------------------------------------------------\n",
       "backend: cuda\n",
       "nbytes: 280.0 MB\n",
       "type: 10000000 * var * float64</pre>"
      ],
      "text/plain": [
       "<Array [[0.5446351700822585, ...], ..., [...]] type='10000000 * var * float64'>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# randomly generate a ListOffsetArray\n",
    "array_size = 10000000\n",
    "# inner arrays lengths \n",
    "lengths = cp.random.randint(0, 6, array_size)\n",
    "# a flat array of all the data\n",
    "layout = ak.contents.NumpyArray(cp.random.random(int(sum(lengths))))\n",
    "# calculate offsets from the lenghts\n",
    "offsets = ak.index.Index(cp.concatenate((cp.array([0]), cp.cumsum(lengths))))\n",
    "\n",
    "rand_arr = ak.Array(ak.contents.ListOffsetArray(offsets, layout))\n",
    "rand_arr"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c95683ce-2af8-4a50-9750-a82f5775ea7b",
   "metadata": {},
   "source": [
    "Save the arrays to the file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "714c46a6-f9b2-425c-a413-178f563df07d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<pyarrow._parquet.FileMetaData object at 0x71eb90c3e700>\n",
       "  created_by: parquet-cpp-arrow version 22.0.0\n",
       "  num_columns: 1\n",
       "  num_rows: 10000000\n",
       "  num_row_groups: 1\n",
       "  format_version: 2.6\n",
       "  serialized_size: 0"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ak.to_parquet(rand_arr, \"random_listoffset_small.parquet\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d6cf5ef5-1fd3-4152-a487-dfceda258cfb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<pyarrow._parquet.FileMetaData object at 0x71ea9049f970>\n",
       "  created_by: parquet-cpp-arrow version 22.0.0\n",
       "  num_columns: 1\n",
       "  num_rows: 90000000\n",
       "  num_row_groups: 2\n",
       "  format_version: 2.6\n",
       "  serialized_size: 0"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ak.to_parquet(ak.concatenate([rand_arr] * 9), \"random_listoffset.parquet\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecf5835f-8eba-420e-a706-d6c4e25e6d71",
   "metadata": {},
   "source": [
    "### Test cccl on bigger data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1b2b8b4d-dbf9-46f8-b3ff-76f9f327d129",
   "metadata": {},
   "outputs": [],
   "source": [
    "awkward_array = ak.to_backend(ak.from_parquet(\"random_listoffset_small.parquet\"), 'cuda')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d43efbd-4cee-410e-bc34-04a55cc13eed",
   "metadata": {},
   "source": [
    "Let's modify our sum function for float64 values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "616854c3-69d6-4018-9689-f7524c9775e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def cccl_sum(awkward_array):\n",
    "    input_data = awkward_array.layout.content.data \n",
    "    offsets = awkward_array.layout.offsets.data\n",
    "    \n",
    "    # Prepare the start and end offsets\n",
    "    start_o = offsets[:-1]\n",
    "    end_o = offsets[1:]\n",
    "    \n",
    "    # Prepare the output array\n",
    "    n_segments = start_o.size\n",
    "    output = cp.empty(n_segments, dtype=cp.float64)\n",
    "    \n",
    "    # Initial value for the reduction\n",
    "    h_init = np.array([0], dtype=np.float64)\n",
    "    \n",
    "    # Perform the segmented reduce\n",
    "    segmented_reduce(\n",
    "        input_data, output, start_o, end_o, sum_op, h_init, n_segments\n",
    "    )\n",
    "    return output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aceb66fd-83d6-42f0-a52a-a2b528dec2a4",
   "metadata": {},
   "source": [
    "Check the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2855e912-a7c7-4b30-869d-9b79ab82f7ac",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([2.03369381, 1.00149052, 0.74000209, ..., 2.11200709, 0.56703317,\n",
       "       0.98825442], shape=(10000000,))"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cccl_sum(awkward_array)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42e274cd-2e81-453f-b7a0-552730f5e499",
   "metadata": {},
   "source": [
    "timeit the performance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a8b88fc7-0543-42e2-b302-1a662ac9a76a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "79.5 μs ± 18.1 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -r 7 -n 100\n",
    "cccl_sum(awkward_array)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3dff449-372f-4775-8c94-fb4a4d474509",
   "metadata": {},
   "source": [
    "Load a 9times bigger array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "00653ae3-9071-4891-8a51-aa1fe14d1a7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "awkward_array2 = ak.to_backend(ak.from_parquet(\"random_listoffset.parquet\"), 'cuda')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "434447af-9fe6-4305-b961-5821e50fb16c",
   "metadata": {},
   "source": [
    "prerun"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2941dab3-b363-4505-b717-42defdb8c4a8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1.05136775, 0.        , 0.62357382, ..., 0.9605404 , 0.41038346,\n",
       "       1.98615626], shape=(90000000,))"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cccl_sum(awkward_array2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f6e7fc1-5ac9-44ea-a16e-b9e5b4bbf72e",
   "metadata": {},
   "source": [
    "timeit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "49dd9653-432b-449d-977c-71458dcd7d25",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "99.4 μs ± 59.9 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -r 7 -n 100\n",
    "cccl_sum(awkward_array2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39d006d1-3c31-4e91-bcde-5c26494a290a",
   "metadata": {},
   "source": [
    "It takes about twice as much time"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ead73e47-7a5d-4b9d-a399-840599813dac",
   "metadata": {},
   "source": [
    "### Test cuda-kernel on bigger data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "111e9c2e-bdff-4ddb-bdc8-3849cba29566",
   "metadata": {},
   "outputs": [],
   "source": [
    "awkward_array = ak.to_backend(ak.from_parquet(\"random_listoffset_small.parquet\"), 'cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "1144347f-494c-4669-ba5d-f663684c74b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_parents(awkward_array, array_data):\n",
    "    offsets = awkward_array.layout.offsets.data\n",
    "\n",
    "    # get the 'parents'\n",
    "    # number of all items\n",
    "    N = len(array_data)\n",
    "    # all the possible parents idx\n",
    "    flat_idx = cp.arange(N)\n",
    "    # use searchsorted to map each flat item to its parent's index\n",
    "    parents = cp.searchsorted(offsets[1:], flat_idx, side=\"right\")\n",
    "    return parents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b726cf44-b44e-4ff1-9418-3cae456424c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cupy as cp\n",
    "\n",
    "cuda_kernel = \"\"\"\n",
    "extern \"C\" {\n",
    "    __global__ void awkward_reduce_sum_a(double* toptr, double* fromptr, int* parents, int lenparents, int outlength) {\n",
    "       int thread_id = blockIdx.x * blockDim.x + threadIdx.x;\n",
    "\n",
    "       if (thread_id < outlength) {\n",
    "          toptr[thread_id] = 0.0f;\n",
    "       }\n",
    "    }\n",
    "}\n",
    "extern \"C\" {\n",
    "    __global__ void awkward_reduce_sum_b(double* toptr, double* fromptr, int* parents, int lenparents, int outlength) {\n",
    "       int thread_id = blockIdx.x * blockDim.x + threadIdx.x;\n",
    "       int stride = blockDim.x * gridDim.x;\n",
    "\n",
    "       for (int i = thread_id; i < lenparents; i += stride) {\n",
    "           atomicAdd(&toptr[parents[i]], fromptr[i]);\n",
    "       }       \n",
    "    }\n",
    "}\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3425524a-eb1b-4c31-b5b3-2a883b659266",
   "metadata": {},
   "source": [
    "Compare the results with cccl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "75dba527-77ed-4e45-873e-ba865ae4a2d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def cuda_kernels_sum(awkward_array):\n",
    "    fromptr = awkward_array.layout.content.data\n",
    "    parents = get_parents(awkward_array, fromptr).astype(cp.int32)\n",
    "    \n",
    "    lenparents = len(parents)\n",
    "    outlength = int(cp.max(parents)) + 1\n",
    "    toptr = cp.zeros(outlength, dtype=cp.float64)\n",
    "    \n",
    "    block_size = 256\n",
    "    grid_size = (lenparents + block_size - 1) // block_size\n",
    "    \n",
    "    raw_module = cp.RawModule(code=cuda_kernel)\n",
    "    \n",
    "    awkward_reduce_sum_a = raw_module.get_function('awkward_reduce_sum_a')\n",
    "    awkward_reduce_sum_b = raw_module.get_function('awkward_reduce_sum_b')\n",
    "    \n",
    "    awkward_reduce_sum_a((grid_size,), (block_size,), (toptr, fromptr, parents, lenparents, outlength))\n",
    "    awkward_reduce_sum_b((grid_size,), (block_size,), (toptr, fromptr, parents, lenparents, outlength))\n",
    "    return toptr.get()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "75488546-954e-41f2-9ae5-fb1c3ab18d94",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([2.03369381, 1.00149052, 0.74000209, ..., 2.11200709, 0.56703317,\n",
       "       0.98825442], shape=(10000000,))"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cuda_kernels_sum(awkward_array)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df8fac69-fe8e-49fe-a725-735ea15669ef",
   "metadata": {},
   "source": [
    "Results are the same! Now let's see how long it takes to use cuda kernels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "9d2e29e4-23ac-4571-8922-09a372cb992b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "29.2 ms ± 633 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -r 7 -n 100\n",
    "cuda_kernels_sum(awkward_array)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "caac9cbc-58fb-416e-b7db-72eb4b2e44f4",
   "metadata": {},
   "source": [
    "test on a 9times bigger array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "46b8632d-e53e-44fa-b9ec-1e76be2bcacc",
   "metadata": {},
   "outputs": [],
   "source": [
    "awkward_array2 = ak.to_backend(ak.from_parquet(\"random_listoffset.parquet\"), 'cuda')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a62aed0-37e3-4e57-a090-5c5a67557207",
   "metadata": {},
   "source": [
    "prerun"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "01161adf-521d-46b1-a696-ca4097415986",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1.05136775, 0.        , 0.62357382, ..., 0.9605404 , 0.41038346,\n",
       "       1.98615626], shape=(90000000,))"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cuda_kernels_sum(awkward_array2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ca0cfcc-cc78-4566-93ea-7ca61dbe0d36",
   "metadata": {},
   "source": [
    "timeit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "cb91381d-e3f8-43f9-a070-78b532733ac2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "236 ms ± 4.87 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -r 7 -n 100\n",
    "cuda_kernels_sum(awkward_array2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73f8f6a9-51e0-485a-9520-9c977f7eb785",
   "metadata": {},
   "source": [
    "This array(~2GB) takes roughly 7 times longer to process then a smaller one(280Mb). Comparing to cccl, it was only a two times longer for a bigger array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58988fda-18db-4a37-8180-2f9cde8a561a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py313",
   "language": "python",
   "name": "py313"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
